package com.djf.springaidemo.memory;

import com.djf.springaidemo.entity.ChatMessage;
import com.djf.springaidemo.mapper.ChatMessageMapper;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;

import java.util.Date;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @Author 26393
 * @Date 2025/5/3 2:06
 * @Version 1.0
 */
@Component
public class MysqlChatMemory implements ChatMemory {

    @Resource
    private ChatMessageMapper chatMessageMapper;

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<ChatMessage> chatMessages = messages.stream()
                .map(message -> {
                    ChatMessage chatMessage = new ChatMessage();
                    chatMessage.setSessionId(conversationId);
                    chatMessage.setContent(message.getText());
                    chatMessage.setRole(message.getMessageType().getValue());
                    chatMessage.setCreatedAt(new Date());
                    return chatMessage;
                })
                .collect(Collectors.toList());
        chatMessageMapper.insert(chatMessages);
    }

    @Override
    public void add(String conversationId, Message message) {
        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setSessionId(conversationId);
        chatMessage.setContent(message.getText());
        chatMessage.setRole(message.getMessageType().getValue());
        chatMessage.setCreatedAt(new Date());
        chatMessageMapper.insert(chatMessage);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        return chatMessageMapper.selectBySessionIdOrderByCreatedAt(conversationId)
                .stream()
                .map(msg -> {
                    switch (msg.getRole()) {
                        case "user":
                            return new UserMessage(msg.getContent());
                        case "system":
                            return new SystemMessage(msg.getContent());
                        default:
                            return new AssistantMessage(msg.getContent());
                    }
                })
                .collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        chatMessageMapper.deleteBySessionId(conversationId);
    }
}
